import torch
import torch.nn as nn

# 模型定义，与训练代码中的模型定义一致
class Perceptron(nn.Module):
    def __init__(self):
        super(Perceptron, self).__init__()
        self.fc = nn.Linear(1, 1)

    def forward(self, x):
        return self.fc(x)

class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(1, 64)
        self.relu = nn.ReLU()  # 明确定义ReLU激活函数
        self.fc2 = nn.Linear(64, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        return self.fc2(x)

# 加载模型参数
def load_model(model_class, model_path):
    model = model_class()
    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
    model.eval()  # 将模型设置为评估模式
    return model

# 推理数据
X_test = torch.tensor([[10.0]], dtype=torch.float32)

# 载入单层感知机模型
model_perceptron = load_model(Perceptron, 'model_perceptron.pth')
with torch.no_grad():
    output_perceptron = model_perceptron(X_test).item()

# 载入多层感知机模型
model_mlp = load_model(MLP, 'model_mlp.pth')
with torch.no_grad():
    output_mlp = model_mlp(X_test).item()

# 打印输出
print("Perceptron Model Structure:", Perceptron)
print(f"Test Input: {X_test.item()}, Test Output (Perceptron): {output_perceptron}")
print("MLP Model Structure:", MLP)
print(f"Test Input: {X_test.item()}, Test Output (MLP): {output_mlp}")